
from typing import List, Dict, Any

from multihop.core.docment import QueryPlan, RetrievalResult, logger
from multihop.core.dpsk_client import DeepSeekClient
from multihop.core.vector_db import VectorDatabase


class MultiHopRAG:
    """多跳RAG检索增强生成系统"""

    def __init__(self, vector_db: VectorDatabase, llm_client: DeepSeekClient, max_hops: int = 3):
        self.vector_db = vector_db
        self.llm_client = llm_client
        self.max_hops = max_hops

    def retrieve_with_hops(self, query_plan: QueryPlan) -> List[RetrievalResult]:
        """执行多跳检索"""
        all_results = []
        seen_doc_ids = set()

        # 初始检索
        initial_results = self.vector_db.search(query_plan.original_query, k=5)
        for result in initial_results:
            if result.document.id not in seen_doc_ids:
                result.hop_level = 0
                all_results.append(result)
                seen_doc_ids.add(result.document.id)
                query_plan.retrieved_docs.append(result)

        # 多跳检索
        for hop in range(1, self.max_hops):
            query_plan.current_hop = hop

            # 构建当前上下文
            context = self._build_context(query_plan.retrieved_docs[-5:])  # 使用最近的5个文档

            # 生成子查询
            sub_queries = self.llm_client.generate_sub_queries(
                query_plan.original_query,
                context
            )

            if not sub_queries:
                logger.info(f"第 {hop} 跳：未生成子查询，结束检索")
                break

            query_plan.sub_queries.extend(sub_queries)

            # 执行子查询检索
            hop_results = []
            for sub_query in sub_queries:
                sub_results = self.vector_db.search(sub_query, k=3)
                for result in sub_results:
                    if result.document.id not in seen_doc_ids:
                        result.hop_level = hop
                        hop_results.append(result)
                        seen_doc_ids.add(result.document.id)

            if not hop_results:
                logger.info(f"第 {hop} 跳：未找到新文档，结束检索")
                break

            # 根据相关性排序并选择top结果
            hop_results.sort(key=lambda x: x.score, reverse=True)
            selected_results = hop_results[:3]  # 每跳最多选择3个文档

            all_results.extend(selected_results)
            query_plan.retrieved_docs.extend(selected_results)

            logger.info(f"第 {hop} 跳：检索到 {len(selected_results)} 个新文档")

        return all_results

    def _build_context(self, results: List[RetrievalResult]) -> str:
        """构建上下文字符串"""
        if not results:
            return ""

        context_parts = []
        for result in results:
            snippet = result.document.content[:200] + "..." if len(
                result.document.content) > 200 else result.document.content
            context_parts.append(f"[跳数{result.hop_level}] {snippet}")

        return "\n\n".join(context_parts)

    def generate_answer(self, query_plan: QueryPlan) -> str:
        """基于检索结果生成最终答案"""
        if not query_plan.retrieved_docs:
            return "抱歉，没有找到相关信息来回答您的问题。"

        # 构建完整上下文
        context = self._build_detailed_context(query_plan.retrieved_docs)

        # 构建生成prompt
        prompt = f"""
基于以下检索到的相关信息，回答用户的问题。请确保答案准确、全面且有逻辑性。

用户问题: {query_plan.original_query}

相关信息:
{context}

请基于上述信息提供详细的答案:
"""

        answer = self.llm_client.generate(prompt, max_tokens=512, temperature=0.3)
        return answer

    def _build_detailed_context(self, results: List[RetrievalResult]) -> str:
        """构建详细的上下文"""
        # 按跳数和相关性分组
        hop_groups = {}
        for result in results:
            hop = result.hop_level
            if hop not in hop_groups:
                hop_groups[hop] = []
            hop_groups[hop].append(result)

        context_parts = []
        for hop in sorted(hop_groups.keys()):
            hop_results = sorted(hop_groups[hop], key=lambda x: x.score, reverse=True)
            for i, result in enumerate(hop_results[:5]):  # 每跳最多5个文档
                context_parts.append(
                    f"[文档 {hop}-{i + 1}, 相关性: {result.score:.3f}]\n{result.document.content}\n"
                )

        return "\n".join(context_parts)

    def query(self, question: str) -> Dict[str, Any]:
        """执行完整的多跳RAG查询"""
        logger.info(f"开始处理查询: {question}")

        # 创建查询计划
        query_plan = QueryPlan(question)

        # 执行多跳检索
        results = self.retrieve_with_hops(query_plan)

        # 生成答案
        answer = self.generate_answer(query_plan)

        return {
            "question": question,
            "answer": answer,
            "retrieved_documents": len(results),
            "hops_used": query_plan.current_hop + 1,
            "sub_queries": query_plan.sub_queries,
            "retrieval_details": [
                {
                    "doc_id": r.document.id,
                    "score": r.score,
                    "hop_level": r.hop_level,
                    "content_snippet": r.document.content[:100] + "..."
                }
                for r in results
            ]
        }